{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c3e2458f-d038-4845-93a0-d4ad830f9f90",
   "metadata": {},
   "source": [
    "# LangChain 核心模块学习：Memory\n",
    "\n",
    "大多数LLM应用都具有对话界面。对话的一个重要组成部分是能够引用先前在对话中介绍过的信息。至少，一个对话系统应该能够直接访问一些过去消息的窗口。更复杂的系统将需要拥有一个不断更新的世界模型，使其能够保持关于实体及其关系的信息。\n",
    "\n",
    "我们将存储过去交互信息的能力称为“记忆（Memory）”。\n",
    "\n",
    "LangChain提供了许多用于向应用/系统中添加 Memory 的实用工具。这些工具可以单独使用，也可以无缝地集成到链中。\n",
    "\n",
    "一个记忆系统（Memory System）需要支持两个基本操作：**读取（READ）和写入（WRITE）**。\n",
    "\n",
    "每个链都定义了一些核心执行逻辑，并期望某些输入。其中一些输入直接来自用户，但有些输入可能来自 Memory。\n",
    "\n",
    "在一个典型 Chain 的单次运行中，将与其 Memory System 进行至少两次交互:\n",
    "\n",
    "1. 在接收到初始用户输入之后，在执行核心逻辑之前，链将从其 Memory 中**读取**并扩充用户输入。\n",
    "2. 在执行核心逻辑之后但在返回答案之前，一个链条将把当前运行的输入和输出**写入** Memory ，以便在未来的运行中可以引用它们。\n",
    "\n",
    "![](../images/memory.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9cf0d43-107b-47ae-9e2c-2edaec38c800",
   "metadata": {},
   "source": [
    "## BaseMemory Class 基类\n",
    "\n",
    "类继承关系：\n",
    "\n",
    "```\n",
    "## 适用于简单的语言模型\n",
    "BaseMemory --> BaseChatMemory --> <name>Memory  # Examples: ZepMemory, MotorheadMemory\n",
    "```\n",
    "\n",
    "**代码实现：https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/schema/memory.py**\n",
    "\n",
    "```python\n",
    "# 定义一个名为BaseMemory的基础类\n",
    "class BaseMemory(Serializable, ABC):\n",
    "    \"\"\"用于Chains中的内存的抽象基类。\n",
    "    \n",
    "    这里的内存指的是Chains中的状态。内存可以用来存储关于Chain的过去执行的信息，\n",
    "    并将该信息注入到Chain的未来执行的输入中。例如，对于会话型Chains，内存可以用来\n",
    "    存储会话，并自动将它们添加到未来的模型提示中，以便模型具有必要的上下文来连贯地\n",
    "    响应最新的输入。\"\"\"\n",
    "\n",
    "    # 定义一个名为Config的子类\n",
    "    class Config:\n",
    "        \"\"\"为此pydantic对象配置。\n",
    "    \n",
    "        Pydantic是一个Python库，用于数据验证和设置管理，主要基于Python类型提示。\n",
    "        \"\"\"\n",
    "    \n",
    "        # 允许在pydantic模型中使用任意类型。这通常用于允许复杂的数据类型。\n",
    "        arbitrary_types_allowed = True\n",
    "    \n",
    "    # 下面是一些必须由子类实现的方法：\n",
    "    \n",
    "    # 定义一个属性，它是一个抽象方法。任何从BaseMemory派生的子类都需要实现此方法。\n",
    "    # 此方法应返回该内存类将添加到链输入的字符串键。\n",
    "    @property\n",
    "    @abstractmethod\n",
    "    def memory_variables(self) -> List[str]:\n",
    "        \"\"\"获取此内存类将添加到链输入的字符串键。\"\"\"\n",
    "    \n",
    "    # 定义一个抽象方法。任何从BaseMemory派生的子类都需要实现此方法。\n",
    "    # 此方法基于给定的链输入返回键值对。\n",
    "    @abstractmethod\n",
    "    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:\n",
    "        \"\"\"根据给链的文本输入返回键值对。\"\"\"\n",
    "    \n",
    "    # 定义一个抽象方法。任何从BaseMemory派生的子类都需要实现此方法。\n",
    "    # 此方法将此链运行的上下文保存到内存。\n",
    "    @abstractmethod\n",
    "    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:\n",
    "        \"\"\"保存此链运行的上下文到内存。\"\"\"\n",
    "    \n",
    "    # 定义一个抽象方法。任何从BaseMemory派生的子类都需要实现此方法。\n",
    "    # 此方法清除内存内容。\n",
    "    @abstractmethod\n",
    "    def clear(self) -> None:\n",
    "        \"\"\"清除内存内容。\"\"\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94c413b6-6e07-411f-bad0-eb46db3a313d",
   "metadata": {},
   "source": [
    "## BaseChatMessageHistory Class 基类\n",
    "\n",
    "类继承关系：\n",
    "\n",
    "```\n",
    "## 适用于聊天模型\n",
    "\n",
    "BaseChatMessageHistory --> <name>ChatMessageHistory  # Example: ZepChatMessageHistory\n",
    "```\n",
    "\n",
    "```python\n",
    "# 定义一个名为BaseChatMessageHistory的基础类\n",
    "class BaseChatMessageHistory(ABC):\n",
    "    \"\"\"聊天消息历史记录的抽象基类。\"\"\"\n",
    "\n",
    "    # 在内存中存储的消息列表\n",
    "    messages: List[BaseMessage]\n",
    "\n",
    "    # 定义一个add_user_message方法，它是一个方便的方法，用于将人类消息字符串添加到存储区。\n",
    "    def add_user_message(self, message: str) -> None:\n",
    "        \"\"\"为存储添加一个人类消息字符串的便捷方法。\n",
    "\n",
    "        参数:\n",
    "            message: 人类消息的字符串内容。\n",
    "        \"\"\"\n",
    "        self.add_message(HumanMessage(content=message))\n",
    "\n",
    "    # 定义一个add_ai_message方法，它是一个方便的方法，用于将AI消息字符串添加到存储区。\n",
    "    def add_ai_message(self, message: str) -> None:\n",
    "        \"\"\"为存储添加一个AI消息字符串的便捷方法。\n",
    "\n",
    "        参数:\n",
    "            message: AI消息的字符串内容。\n",
    "        \"\"\"\n",
    "        self.add_message(AIMessage(content=message))\n",
    "\n",
    "    # 抽象方法，需要由继承此基类的子类来实现。\n",
    "    @abstractmethod\n",
    "    def add_message(self, message: BaseMessage) -> None:\n",
    "        \"\"\"将Message对象添加到存储区。\n",
    "\n",
    "        参数:\n",
    "            message: 要存储的BaseMessage对象。\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    # 抽象方法，需要由继承此基类的子类来实现。\n",
    "    @abstractmethod\n",
    "    def clear(self) -> None:\n",
    "        \"\"\"从存储中删除所有消息\"\"\"\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e716cd45-1a71-44da-a924-09d6a56ff6c8",
   "metadata": {},
   "source": [
    "### ConversationChain and ConversationBufferMemory\n",
    "\n",
    "`ConversationBufferMemory` 可以用来存储消息，并将消息提取到一个变量中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d6b700e0-abd4-4531-ad93-b278357d9c64",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import OpenAI\n",
    "from langchain.chains import ConversationChain\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "\n",
    "llm = OpenAI(temperature=0)\n",
    "conversation = ConversationChain(\n",
    "    llm=llm, \n",
    "    verbose=True, \n",
    "    memory=ConversationBufferMemory()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5d51fbb4-1d8e-4ec1-8c55-ec70247d4d64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi there!\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi there! It's nice to meet you. How can I help you today?\""
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Hi there!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ac5411e7-b8ec-4c31-b659-deb44af038df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there!\n",
      "AI:  Hi there! It's nice to meet you. How can I help you today?\n",
      "Human: I'm doing well! Just having a conversation with an AI.\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" That's great! It's always nice to have a conversation with someone new. What would you like to talk about?\""
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"I'm doing well! Just having a conversation with an AI.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7428157f-72ed-4b8d-b114-80bfb96e13bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi there!\n",
      "AI:  Hi there! It's nice to meet you. How can I help you today?\n",
      "Human: I'm doing well! Just having a conversation with an AI.\n",
      "AI:  That's great! It's always nice to have a conversation with someone new. What would you like to talk about?\n",
      "Human: Tell me about yourself.\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Sure! I'm an AI created to help people with their everyday tasks. I'm programmed to understand natural language and provide helpful information. I'm also constantly learning and updating my knowledge base so I can provide more accurate and helpful answers.\""
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation.predict(input=\"Tell me about yourself.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f66476e0-6d20-4ba1-98af-fd5300096a8c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "694c1e3e-4024-4cc3-963e-01fe1a60f1c3",
   "metadata": {},
   "source": [
    "### ConversationBufferWindowMemory\n",
    "`ConversationBufferWindowMemory` 会在时间轴上保留对话的交互列表。它只使用最后 K 次交互。这对于保持最近交互的滑动窗口非常有用，以避免缓冲区过大。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "dbe61067-5b8f-40a1-827e-4c8c6ad473dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "\n",
      "Human: Hi, what's up?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi there! I'm doing great. I'm currently helping a customer with a technical issue. How about you?\""
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.memory import ConversationBufferWindowMemory\n",
    "\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=OpenAI(temperature=0), \n",
    "    # We set a low k=2, to only keep the last 2 interactions in memory\n",
    "    memory=ConversationBufferWindowMemory(k=2), \n",
    "    verbose=True\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, what's up?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ae850ece-78b8-41ad-97ea-91a57a275a8b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi, what's up?\n",
      "AI:  Hi there! I'm doing great. I'm currently helping a customer with a technical issue. How about you?\n",
      "Human: What's their issues?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" The customer is having trouble connecting to their Wi-Fi network. I'm helping them troubleshoot the issue and get them connected.\""
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"What's their issues?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9e2fe36d-44ec-4351-8922-4481c2bf6750",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi, what's up?\n",
      "AI:  Hi there! I'm doing great. I'm currently helping a customer with a technical issue. How about you?\n",
      "Human: What's their issues?\n",
      "AI:  The customer is having trouble connecting to their Wi-Fi network. I'm helping them troubleshoot the issue and get them connected.\n",
      "Human: Is it going well?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Yes, it's going well so far. We've already identified the problem and are now working on a solution.\""
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conversation_with_summary.predict(input=\"Is it going well?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1db201fd-1373-4148-ab04-525ea089a9fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Current conversation:\n",
      "Human: What's their issues?\n",
      "AI:  The customer is having trouble connecting to their Wi-Fi network. I'm helping them troubleshoot the issue and get them connected.\n",
      "Human: Is it going well?\n",
      "AI:  Yes, it's going well so far. We've already identified the problem and are now working on a solution.\n",
      "Human: What's the solution?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" The solution is to reset the router and reconfigure the settings. We're currently in the process of doing that.\""
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Notice here that the first interaction does not appear.\n",
    "conversation_with_summary.predict(input=\"What's the solution?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fc35065-ff20-4fda-ac5b-0976102160a9",
   "metadata": {},
   "source": [
    "### ConversationSummaryBufferMemory\n",
    "\n",
    "`ConversationSummaryBufferMemory` 在内存中保留了最近的交互缓冲区，但不仅仅是完全清除旧的交互，而是将它们编译成摘要并同时使用。与以前的实现不同的是，它使用标记长度而不是交互次数来确定何时清除交互。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "174ccc64-2dd9-4c98-b638-6aa542bdbd55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationSummaryBufferMemory\n",
    "\n",
    "memory = ConversationSummaryBufferMemory(llm=llm, max_token_limit=10)\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c520d4eb-f095-4245-b003-83581b619b2b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'System: \\nThe human greets the AI, to which the AI responds.\\nHuman: not much you\\nAI: not much'}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "909f7407-eb70-423a-bb81-a77b48875280",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = ConversationSummaryBufferMemory(\n",
    "    llm=llm, max_token_limit=10, return_messages=True\n",
    ")\n",
    "memory.save_context({\"input\": \"hi\"}, {\"output\": \"whats up\"})\n",
    "memory.save_context({\"input\": \"not much you\"}, {\"output\": \"not much\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "849e3c77-790c-431c-a20e-f4aa7f809335",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nThe human and AI exchange that they are not doing much.'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = memory.chat_memory.messages\n",
    "previous_summary = \"\"\n",
    "memory.predict_new_summary(messages, previous_summary)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81c7b148-5988-4427-90da-ab7bdc4870b6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
